from __future__ import annotations

import pathlib, os, subprocess
from typing import Mapping
from google.protobuf.compiler import plugin_pb2 as plugin
from google.protobuf.descriptor import Descriptor, EnumDescriptor
from google.protobuf.descriptor_pool import DescriptorPool
from google.protobuf.descriptor_pb2 import (
    DescriptorProto,
    EnumDescriptorProto,
    FieldDescriptorProto,
    FileDescriptorProto,
    ServiceDescriptorProto,
    OneofDescriptorProto,
    FileDescriptorSet,
)

data = b"""
syntax = "proto3";

package example;

message Person {
    message Contact {
        string qq = 1;
        string wechat = 2;
    }
    enum Gender {
        MALE = 0;
        FEMALE = 1;
    }
    string name = 1;
    Contact contact = 2;
}
"""


def find_message_by_name(pool: DescriptorPool, full_name: str) -> Descriptor | None:
    try:
        d = pool.FindMessageTypeByName(full_name)
        print(d.full_name)
        return d
    except KeyError as e:
        print("Error")
        return None


def find_enum_by_name(pool: DescriptorPool, full_name: str) -> EnumDescriptor | None:
    try:
        d = pool.FindEnumTypeByName(full_name)
        print(d.full_name)
        return d
    except KeyError as e:
        print("Error")
        return None


def main():
    temp_proto_base = "/tmp/proto"
    os.makedirs(temp_proto_base, exist_ok=True)
    temp_proto = f"{temp_proto_base}/temp.proto"
    temp_pd = f"{temp_proto_base}/temp.pd"
    pathlib.Path(temp_proto).write_bytes(data)
    subprocess.call(
        f"protoc --descriptor_set_out={temp_pd} --include_imports -I {temp_proto_base} {temp_proto}".split()
    )
    fds = FileDescriptorSet.FromString(pathlib.Path(temp_pd).read_bytes())
    pool = DescriptorPool()
    for fd in fds.file:
        print(f">>> Adding {fd.name}")
        pool.Add(fd)
    find_message_by_name(pool, "example.Person")
    find_message_by_name(pool, "example.Person.Contact")
    find_enum_by_name(pool, "example.Person.Gender")
    find_message_by_name(pool, ".example.Person")
    find_message_by_name(pool, ".example.Person.Contact")
    find_enum_by_name(pool, ".example.Person.Gender")


if __name__ == "__main__":
    main()
